package loan.calc.datasource.mytransaction;

import loan.calc.datasource.DynamicDataSource;
import lombok.extern.slf4j.Slf4j;
import org.apache.ibatis.transaction.Transaction;
import org.springframework.jdbc.datasource.DataSourceUtils;

import javax.sql.DataSource;
import java.sql.Connection;
import java.sql.SQLException;
import java.util.concurrent.ConcurrentHashMap;

/**
 * @author v_ECD963
 */
@Slf4j
public class MultiDataSourceTransaction implements Transaction {
    private final DynamicDataSource dynamicDataSource;
    private final ConcurrentHashMap<String, DataSource> dataSourceMap;
    private final ConcurrentHashMap<String, Connection> connectionMap;
    private final ConcurrentHashMap<String, Boolean> autoCommitMap;
    private final ConcurrentHashMap<String, Boolean> isConnectionTransactionalMap;

    public MultiDataSourceTransaction(DataSource dataSource) {
        this.dynamicDataSource = (DynamicDataSource) dataSource;
        dataSourceMap = new ConcurrentHashMap<>();
        connectionMap = new ConcurrentHashMap<>();
        autoCommitMap = new ConcurrentHashMap<>();
        isConnectionTransactionalMap = new ConcurrentHashMap<>();
    }

    public Connection getConnection() throws SQLException {
        String dsKey = DynamicDataSource.getCurrent();
        if (!dataSourceMap.containsKey(dsKey)) {
            DataSource dataSource = dynamicDataSource.getAcuallyDataSource();
            dataSourceMap.put(dsKey, dataSource);
        }
        if (!connectionMap.containsKey(dsKey)) {
            Connection connection = DataSourceUtils.getConnection(dataSourceMap.get(dsKey));
            connectionMap.put(dsKey, connection);
        }
        if (!autoCommitMap.containsKey(dsKey)) {
            boolean autoCommit = connectionMap.get(dsKey).getAutoCommit();
            autoCommitMap.put(dsKey, autoCommit);
        }
        if (!isConnectionTransactionalMap.containsKey(dsKey)) {
            boolean isConnectionTransactional = DataSourceUtils.isConnectionTransactional(connectionMap.get(dsKey), dataSourceMap.get(dsKey));
            isConnectionTransactionalMap.put(dsKey, isConnectionTransactional);
        }
        return connectionMap.get(dsKey);
    }


    public void commit() throws SQLException {
        for (String dsKey : connectionMap.keySet()) {
            Connection connection = connectionMap.get(dsKey);
            boolean isConnectionTransactional = isConnectionTransactionalMap.get(dsKey);
            boolean autoCommit = autoCommitMap.get(dsKey);
            if (connection != null && !isConnectionTransactional && !autoCommit) {
                connection.commit();
            }
        }
    }

    public void rollback() throws SQLException {
        for (String dsKey : connectionMap.keySet()) {
            Connection connection = connectionMap.get(dsKey);
            boolean isConnectionTransactional = isConnectionTransactionalMap.get(dsKey);
            boolean autoCommit = autoCommitMap.get(dsKey);
            if (connection != null && !isConnectionTransactional && !autoCommit) {
                connection.rollback();
            }
        }
    }

    public void close() {
        for (String dsKey : connectionMap.keySet()) {
            Connection connection = connectionMap.get(dsKey);
            DataSource dataSource = dataSourceMap.get(dsKey);
            DataSourceUtils.releaseConnection(connection, dataSource);
        }
    }

    public Integer getTimeout() {
        return null;
    }
}
